import sqlite3
import tempfile

import pandas as pd
import streamlit as st
from langchain_community.utilities import SQLDatabase

from langchain_db.db_assitant import get_response


# def create_chart(input_data, chart_type):
#     df_data = pd.DataFrame(input_data["data"], columns=input_data["columns"])
#     df_data.set_index(input_data["columns"][0], inplace=True)
#     if chart_type == "table":
#         st.table(df_data)
#     elif chart_type == "bar":
#         st.bar_chart(df_data)
#     elif chart_type == "line":
#         st.line_chart(df_data)
#     elif chart_type == "scatter":
#         st.line_chart(df_data)


with st.sidebar:
    openai_key = st.text_input("输入OpenAI API Key", type="password")

st.header("SQL智能数据分析")
db_file = st.file_uploader("上传数据库文件", type=["db"])
if db_file:
    tfile = tempfile.NamedTemporaryFile(delete=False)
    tfile.write(db_file.getvalue())
    st.session_state["db_file"] = tfile.name
    with st.expander("展示数据"):
        conn = sqlite3.connect(st.session_state["db_file"])
        cursor = conn.cursor()
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
        tables = cursor.fetchall()
        st.write("数据库中的表：", tables)

question = st.text_input("请根据上面提交的数据库文件，进行问题提问")
flag = st.button("提交问题")

if flag:
    if 'db_file' not in st.session_state:
        st.warning("请先上传数据库文件")
        st.stop()
    if not question:
        st.warning("请输入问题")
        st.stop()
    if not openai_key:
        st.info("请输入你的Openai密钥")
        st.stop()
    with st.spinner("正在分析数据"):
        db_session = st.session_state["db_file"]
        db = SQLDatabase.from_uri(f"sqlite:///{db_session}")
        response = get_response(openai_key, question, db)
        st.write(response["output"])
        # if "table" in response:
        #     create_chart(response["table"], "table")
        # if "bar" in response:
        #     create_chart(response["bar"], "bar")
        # if "line" in response:
        #     create_chart(response["line"], "line")
        # if "scatter" in response:
        #     create_chart(response["scatter"], "scatter")
